package wntm;

import edu.cn.nlsde.Utils.SparseMatrix;

import java.io.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Prepare input data for Word Network Topic Model(WNTM)
 *
 * @author Yuan Zuo
 */
public class PrepareInput_co {
    private SparseMatrix word_coOccur = new SparseMatrix();

    // Store original documents, store word IDs to reduce memory consumption
    private List<List<Integer>> docs = new ArrayList<List<Integer>>();

    // Dictionary: word map to its ID and ID map to word
    private Map<String, Integer> w2i = new HashMap<String, Integer>();
    private Map<Integer, String> i2w = new HashMap<Integer, String>();

    // Node map to its adjacent list
    private Map<Integer, Map<Integer, Integer>> n2e = new HashMap<Integer, Map<Integer, Integer>>();

    // Please refer to paper for the definitions of a node's degree and activity
    // Nodes' degrees: node(word id) map to its degree
    private Map<Integer, Integer> n2d = new HashMap<Integer, Integer>();

    // Nodes' activity: node(word id) map to its activity
    private Map<Integer, Integer> n2a = new HashMap<Integer, Integer>();

    public static void main(String[] args) {
        PrepareInput_co preIput = new PrepareInput_co();
        String root = "data/20newsgroup/wntm_co/";
        String orgFile = root + "20newsgroup.data";
        String outDir = root;
        String pseName = "20newsgroup";
        int windowSize = 10;
        preIput.prepareInput(orgFile, outDir + pseName, windowSize);

    }

    private BufferedReader getReader(String path, String charset) {
        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new InputStreamReader(
                    new FileInputStream(path), charset));
        } catch (IOException e) {
            e.printStackTrace();
        }
        return reader;
    }

    private BufferedWriter getWriter(String path, String charset) {
        BufferedWriter writer = null;
        try {
            writer = new BufferedWriter(new OutputStreamWriter(
                    new FileOutputStream(path), charset));
        } catch (IOException e) {
            e.printStackTrace();
        }
        return writer;
    }

    private void loadTxt(String org) {
        BufferedReader reader = getReader(org, "utf-8");

        String line;

        System.out.println("Loading txt file ...");

        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();
                String[] tokens = line.trim().split("\\s+");

                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }

                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private int sumOfAdjacentWeights(Map<Integer, Integer> map) {
        int sum = 0;
        for (int key : map.keySet())
            sum += map.get(key);
        return sum;
    }

    /**
     * @param ws: Sliding window size
     */
    private void buildWordNet(int ws) {
        System.out.println("Building Word Network ...");
        int D = docs.size();
        int prePercent = 0;
        for (int d = 0; d != D; d++) {

            if (d % 1000 == 0) {
                int percent = (int) Math.ceil(100 * ((double) d / (double) D));
                if (prePercent != percent) {
                    prePercent = percent;
                    System.out.println(String.format("%d%%", percent));
                }
            }

            int N = docs.get(d).size();
            int totalSlides = N <= ws ? 1 : (N - ws + 1);

            for (int s = 0; s != totalSlides; s++) {
                int toIndex = N <= ws ? N : (s + ws);
                List<Integer> window = docs.get(d).subList(s, toIndex);
                for (Integer node1 : window) {
                    for (Integer node2 : window) {
                        if (node1.compareTo(node2) != 0 && word_coOccur.get(node1, node2) > 1) {
                            // Update n2e
                            if (n2e.containsKey(node1)) {
                                Map<Integer, Integer> n1Edges = n2e.get(node1);
                                if (n1Edges.containsKey(node2))
                                    n1Edges.put(node2, n1Edges.get(node2) + 1);
                                else
                                    n1Edges.put(node2, 1);
                            } else {
                                Map<Integer, Integer> n1Edges = new HashMap<Integer, Integer>();
                                n1Edges.put(node2, 1);
                                n2e.put(node1, n1Edges);
                            }
                        }
                    }
                }
            }
        }

        // Compute n2d and n2a
        for (int node : n2e.keySet()) {
            int degree = sumOfAdjacentWeights(n2e.get(node));
            int activity = (int) Math.ceil((double) degree
                    / (double) n2e.get(node).size());
            n2d.put(node, degree);
            n2a.put(node, activity);
        }
    }

    private void networkReweight() {

        System.out.println("Network Re-weighting ...");

        for (int node1 : n2e.keySet()) {
            Map<Integer, Integer> n1Edges = n2e.get(node1);
            for (int node2 : n2e.get(node1).keySet()) {
                int tnode = n2d.get(node1) <= n2d.get(node2) ? node1 : node2;
                int activity = n2a.get(tnode);
                n1Edges.put(node2, (int) Math.ceil((double) n1Edges.get(node2) / (double) activity));
            }
            n2e.put(node1, n1Edges);
        }
    }

    private void writeOutput(String pseudoName) {

        System.out.println("Writeing output ...");

        BufferedWriter wWriter = getWriter(pseudoName + ".words", "utf-8");
        BufferedWriter aWriter = getWriter(pseudoName + ".adjacent", "utf-8");

        try {
            //aWriter.append(n2e.size()+"\n");

            for (int node1 : n2e.keySet()) {
                Map<Integer, Integer> n1Edges = n2e.get(node1);
                wWriter.append(i2w.get(node1) + "\n");
                for (int node2 : n1Edges.keySet()) {
                    for (int i = 0; i != n1Edges.get(node2); i++) {
                        aWriter.append(i2w.get(node2) + " ");
                    }
                }
                aWriter.append("\n");
            }
            wWriter.flush();
            wWriter.close();
            aWriter.flush();
            aWriter.close();
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    /**
     * Build word network from original corpus and re-weighting the network, then output pseudo documents.
     *
     * @param org:        Original corpus
     * @param pseudoName: Pseudo documents i.e. input for WNTM
     */
    public void prepareInput(String org, String pseudoName, int ws) {
        loadTxt(org);
        get_word_co_occur();
        buildWordNet(ws);
        networkReweight();
        writeOutput(pseudoName);
    }

    private void get_word_co_occur() {
        for (List<Integer> doc : docs) {
            int length = doc.size();
            for (int i = 0; i < length - 1; i++) {
                for (int j = i + 1; j < length; j++) {
                    int w1 = doc.get(i);
                    int w2 = doc.get(j);
                    if (w1 != w2) {
                        word_coOccur.set(w1, w2, word_coOccur.get(w1, w2) + 1);
                        word_coOccur.set(w2, w1, word_coOccur.get(w2, w1) + 1);
                    }
                }
            }
        }

    }
}
